import chinese_calendar

import pandas as pd
import datetime as dt

import threading
import os
import traceback

from formula import *
from sendmsg import *
from tdx import *
from bars import *


def _isStockMktOpen() -> bool:
    """当前时间是否股市开市中"""
    MORNING_START = dt.time(9, 29, 0)
    MORNING_END = dt.time(11, 30, 0)
    AFTERNOON_START = dt.time(13, 29, 0)
    AFTERNOON_END = dt.time(14, 59, 0)

    time_cur = dt.datetime.now().time()
    if (time_cur > MORNING_START and time_cur < MORNING_END) or \
            (time_cur > AFTERNOON_START and time_cur < AFTERNOON_END):
        return True
    else:
        return False


def _isTimePoint(period=30) -> bool:
    """判断当前时间点是否是周期监控时间点
    输入：period=5/15/30/60/120/240
    """
    TIME_POINT = {'60': [(10, 30), (11, 30), (14, 0)], '120': [(11, 30)], '240': []}

    minute_cur = dt.datetime.now().time().minute + 1
    hour_cur = dt.datetime.now().time().hour

    if (hour_cur, minute_cur) == (9, 30):  # 开市时做一次判断
        return True
    elif (hour_cur, minute_cur) == (14, 57):  # 收市时用14:56代替15:00的判断
        return True
    elif period < 60:  # period = 5/15/30
        if minute_cur % period == 0:
            return True
    elif (hour_cur, minute_cur) in TIME_POINT[str(period)]:  # period = 60/120/240
        return True

    return False


class StockMonitor():
    """股票市场监控器
    按照设置的实时监控、周期监控规则对指定股票的行情进行告警
    """

    def __init__(self, interval=60, log_f=print, alert_f=1):

        # 定义规则集
        self.period_rules = pd.DataFrame(columns=['stock', 'MA', 'period'])
        self.realtime_rules = pd.DataFrame(columns=['stock', 'lower', 'upper', 'rate', 'alert_times'])
        self.trend_rules = pd.DataFrame(columns=['stock', 'MA', 'period', 'rate', 'offset'])

        self.interval = interval  # 秒
        self.log_f = log_f
        self.log_f('Start StockMonitor: time interval=%ds.' % interval)

        self.alert_f = alert_f  # 0-本地print，1-wechat msg
        if self.alert_f == 1:
            # 测试一下robot
            if not send_msg('Stock monitor starts at '+str(dt.datetime.now())):
                print("Can't send message to Wechat !")
            else:
                print("Connected to Wechat robot.")
                
    def set_interval(self, interval=1) -> bool:
        """设定监控时间间隔
        输入：interval，单位：秒，最小为1，最大为60
        注意：如果>60会影响period 监控
        """
        if interval < 1 or interval > 60:
            self.log_f(
                'StockMonitor time interval should be between 1 and 60 seconds.', 'ERROR')
            return False
        self.log_f('Change StockMonitor time interval(s): %d -> %d.' %
                   (self.interval, interval))
        self.interval = interval
        self.run()

    def add_real_time_rule(self, stock='', lower=0.0, upper=0.0, rate=1.0) -> bool:
        """添加实时监控
        """
        if any(self.realtime_rules['stock'] == stock):
            return False
        else:
            self.realtime_rules = self.realtime_rules.append(pd.DataFrame([[stock, lower, upper, rate, 0]], columns=['stock', 'lower', 'upper', 'rate', 'alert_times']), ignore_index=True)
            self.log_f('Add rule: %s,(%.3f,%.3f),%.1f' % (stock, lower, upper, rate))
            return True

    def delete_real_time_rule(self, idx=0):
        """删除指定序号的实时监控
        """
        l = self.realtime_rules.iloc[idx].to_list()
        self.log_f('Delete rule: {0},({1},{2}),{3},{4}'.format(*l))
        self.realtime_rules.drop(index=idx, inplace=True)
        self.realtime_rules.reset_index(drop=True, inplace=True)

    def _isExist(self, rules, stock, *args) -> bool:
        """判断MA规则是否存在"""
        rd = rules.loc[rules['stock'] == stock]
        if rd.empty:
            return False
        for idx, row in rd.iterrows():
            if len(args) == 2 and (row['MA'], row['period']) == (args[0], args[1]):
                return True
            elif len(args) == 3 and (row['MA'], row['period'], row['offset']) == (args[0], args[1], args[2]):
                return True
        return False

    def add_period_rule(self, stock='', ma=5, period=30) -> bool:
        """添加周期监控规则
        输入：ma=5/10/20/30/60 period=30/60/240
        """
        if self._isExist(self.period_rules, stock, ma, period):
            return False
        self.period_rules = self.period_rules.append(pd.DataFrame([[stock, ma, period]], columns=[
            'stock', 'MA', 'period']), ignore_index=True)
        self.log_f('Add rule: %s,MA%d,%dm' % (stock, ma, period))
        return True

    def delete_period_rule(self, idx=0):
        """删除周期监控规则"""
        l = self.period_rules.iloc[idx].to_list()
        self.log_f('Delete rule: {0},MA{1},{2}m'.format(*l))
        self.period_rules.drop(index=idx, inplace=True)
        self.period_rules.reset_index(drop=True, inplace=True)

    def add_trend_rule(self, stock='', ma=10, period=120, rate=0.0, offset=1) -> bool:
        """添加趋势监控规则
        输入：ma=5/10/20/30 period=120/240 rate=(0.0, 5.0) offset=1/2/3
        """
        if self._isExist(self.trend_rules, stock, ma, period, offset):
            return False
        self.trend_rules = self.trend_rules.append(pd.DataFrame([[stock, ma, period, rate, offset]], columns=[
            'stock', 'MA', 'period', 'rate', 'offset']), ignore_index=True)
        self.log_f('Add rule: %s,MA%d,%dm,%.1f,%d' % (stock, ma, period, rate, offset))
        return True

    def delete_trend_rule(self, idx=0):
        """删除趋势监控规则"""
        l = self.trend_rules.iloc[idx].to_list()
        self.log_f('Delete rule: {0},MA{1},{2}m,{3},{4}'.format(*l))
        self.trend_rules.drop(index=idx, inplace=True)
        self.trend_rules.reset_index(drop=True, inplace=True)

    def run(self):
        """启动定时器，按interval处理所有规则
        """
        if chinese_calendar.is_holiday(dt.datetime.now().date()):
            self.log_f(
                "Today is holiday, market is not open!", "WARNING")
            return

        if _isStockMktOpen():  # 在股市开盘时间内

            self.log_f('Start checking all the rules.')

            # 遍历检查所有实时监控规则
            for i in range(0, len(self.realtime_rules)):
                try:
                    stock, lower, upper, rate, alert_times = self.realtime_rules.iloc[i].tolist()
                    self.log_f('Checking rule: %s,(%.3f-%.3f),%.1f%%,%d' % (stock, lower, upper, rate, alert_times))

                    price_cur = getBars(stock, '1', 1).at[0, 'close']
                    # wechat alert
                    if alert_times == 0:
                        if price_cur > upper:
                            self.alert('%s %.3f 高于价格区间上限%.3f' %
                                       (stock, price_cur, upper))
                            self.realtime_rules.at[i, 'alert_times'] = 1
                        elif price_cur < lower:
                            self.alert('%s %.3f 低于价格区间下限%.3f' %
                                       (stock, price_cur, lower))
                            self.realtime_rules.at[i, 'alert_times'] = 1
                    elif alert_times == 1:
                        if price_cur > upper*(1.0+rate/100.0):
                            self.alert('%s %.3f 高于价格区间上限%.3f外扩%.1f%%' %
                                       (stock, price_cur, upper, rate))
                            self.realtime_rules.at[i, 'alert_times'] = 2
                        elif price_cur < lower*(1.0-rate/100.0):
                            self.alert('%s %.3f 低于价格区间下限%.3f外扩%.1f%%' %
                                       (stock, price_cur, lower, rate))
                            self.realtime_rules.at[i, 'alert_times'] = 2
                except:
                    traceback.print_exc()
                    self.log_f("Can't get K data from TDX server!", 'ERROR')

            # 遍历检查所有周期性规则
            for i in range(0, len(self.period_rules)):
                try:
                    stock, ma, period = self.period_rules.iloc[i].tolist()

                    # 判断当前时间点是否是周期监控时间点
                    if not _isTimePoint(period):
                        continue
                    else:
                        self.log_f('Checking rule: %s,%dm,MA%d' % (stock, period, ma))

                    rd = getBars(stock, str(period), ma+1)

                    ma_pre = pMA(rd['close'], ma-1, ma)
                    price_pre = rd.at[ma-1, 'close']
                    ma_cur = cMA(rd['close'], ma, ma)
                    price_cur = rd.at[ma, 'close']

                    # print(stock, ma, period)
                    # print(rd)
                    # print(ma_cur, price_cur, ma_pre, price_pre)

                    if price_pre < ma_pre and price_cur > ma_cur:  # 上穿
                        self.alert('%s %.3f %d分钟级别，上穿%d均' %
                                   (stock, price_cur, period, ma))
                    elif price_pre > ma_pre and price_cur < ma_cur:  # 跌破
                        self.alert('%s %.3f %d分钟级别，跌破%d均' %
                                   (stock, price_cur, period, ma))
                except:
                    traceback.print_exc()
                    self.log_f("Can't get K data from TDX server!", 'ERROR')

            # 遍历检查所有趋势规则
            for i in range(0, len(self.trend_rules)):
                try:
                    stock, ma, period, rate, offset = self.trend_rules.iloc[i].tolist()
                    self.log_f('Checking rule: %s,%dm,MA%d,%.1f,%d' % (stock, period, ma, rate, offset))

                    rd = getBars(stock, str(period), ma+1)
                    rd.drop(index=len(rd)-1, inplace=True)  # 丢弃当前周期的数据

                    ma_pre = pMA(rd['close'], ma-1, ma)
                    ma_cur = cMA(rd['close'], ma, ma)
                    ma_next = ma_cur + (ma_cur - ma_pre)*offset

                    price_cur = getBars(stock, '1', 1).at[0, 'close']
                    price_next = ma_next * (1+rate/100.0)

                    if price_cur <= price_next:
                        self.alert('%s，%d线，逼近%d天的%d均，价%.3f，上浮%.1f%%，价%.3f' % (stock, period, offset, ma, ma_next, rate, price_next))

                except:
                    traceback.print_exc()
                    self.log_f("Can't get K data from TDX server!", 'ERROR')

            self.log_f('Finished checking.')
        else:
            self.log_f("Not the right time: stock market is closed.")

        self.timer = threading.Timer(self.interval, self.run)
        self.timer.start()

    def stop(self):
        self.timer.cancel()

    def alert(self, msg=''):
        """发现异常时的告警函数
        """
        self.log_f(msg, 'ALERT')

        if self.alert_f != 0:
            if not send_msg(msg):
                self.log_f('Can\'t send message !', 'ERROR')

    def save_to_config(self):
        """把所有策略保存到配置文件：实时监控，MA规则
        """
        os.path.exists('.config/') or os.makedirs('.config/')
        self.period_rules.to_csv('.config/period_rules.csv')
        self.realtime_rules.to_csv('.config/realtime_rules.csv')
        self.trend_rules.to_csv('.config/trend_rules.csv')

    def load_from_config(self):
        """从配置文件读取所有策略"""
        realtime_rules_pathfile = '.config/realtime_rules.csv'
        period_rules_pathfile = '.config/period_rules.csv'
        trend_rules_pathfile = '.config/trend_rules.csv'

        if os.path.exists(realtime_rules_pathfile):
            self.realtime_rules = pd.read_csv(realtime_rules_pathfile, index_col=0, dtype={'stock': str})
            self.realtime_rules['alert_times'] = 0  # 重置alert_times
        if os.path.exists(period_rules_pathfile):
            self.period_rules = pd.read_csv(period_rules_pathfile, index_col=0, dtype={'stock': str})
        if os.path.exists(trend_rules_pathfile):
            self.trend_rules = pd.read_csv(trend_rules_pathfile, index_col=0, dtype={'stock': str})


def _test():
    def _logger(msg='', level='INFO'):
        cur_time = dt.datetime.now().strftime('%H:%M:%S ')
        print(cur_time + level+' '+msg)

    sm = StockMonitor(60, _logger, 0)
    # sm.load_from_config()

    # 159928 159995 159825
    sm.add_period_rule('512660', 30, 30)
    sm.add_period_rule('515030', 30, 30)
    sm.add_period_rule('159995', 30, 30)
    sm.add_period_rule('159949', 30, 30)
    sm.add_period_rule('159928', 30, 30)
    sm.add_period_rule('510500', 30, 30)

    #
    sm.add_real_time_rule('159995', 1.6, 1.7, 5.0)
    sm.add_real_time_rule('159928', 1.165, 1.175, 5.0)

    #
    sm.add_trend_rule('159995', 30, 120, 5.0, 2)
    sm.add_trend_rule('159928', 20, 240, 5.0, 2)

    # 启动
    sm.run()
    # sm.save_to_config()


if __name__ == "__main__":
    _test()
